#!/usr/bin/env python3

import rospy
from geometry_msgs.msg import Point
import numpy as np
import cv2

def transform_to_center_origin(image_shape, x, y):
    height, width = image_shape[0], image_shape[1]
    center_x = width // 2
    center_y = height // 2

    dx = int(x) - center_x
    dy = center_y - int(y)

    return dx, dy


def detect_red_block_center(image):
    image_copy = image.copy()
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    lower_red1 = np.array([0, 100, 100])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([170, 100, 100])
    upper_red2 = np.array([180, 255, 255])

    mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask = cv2.bitwise_or(mask_red1, mask_red2)

    mask = cv2.erode(mask, None, iterations=2)
    mask = cv2.dilate(mask, None, iterations=4)
    mask = cv2.erode(mask, None, iterations=2)

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    min_area = 100
    max_area = 5000
    contours = [cnt for cnt in contours if min_area < cv2.contourArea(cnt) < max_area]

    if not contours:
        return None, None, image_copy, None

    max_contour = max(contours, key=cv2.contourArea)
    max_area = cv2.contourArea(max_contour)

    x, y, w, h = cv2.boundingRect(max_contour)
    cx = x + w // 2
    cy = y + h // 2

    cv2.rectangle(image_copy, (x, y), (x + w, y + h), (0, 255, 0), 2)
    cv2.circle(image_copy, (cx, cy), 5, (0, 0, 255), -1)

    return cx, cy, image_copy, max_contour


def red_block_publisher():
    rospy.init_node('red_block_detector', anonymous=True)
    pub = rospy.Publisher('/red_block/offset', Point, queue_size=10)

    rate = rospy.Rate(10)  # 10Hz

    cap = cv2.VideoCapture('/dev/video0')

    while not rospy.is_shutdown():
        ret, frame = cap.read()
        if not ret:
            rospy.logwarn("无法获取视频帧")
            continue

        cx, cy, frame_drawn, max_contour = detect_red_block_center(frame)
        height, width = frame.shape[0], frame.shape[1]
        point_msg = Point()
        #rospy.loginfo(f"当前图像分辨率: {width}x{height}")
        if cx is not None and cy is not None:
            offset_x, offset_y = transform_to_center_origin((height, width), cx, cy)
            real_offset_x = offset_y     # X轴：向上为正
            real_offset_y = -offset_x    # Y轴：向左为正
            area = cv2.contourArea(max_contour)

            point_msg.x = float(real_offset_x)
            point_msg.y = float(real_offset_y)
            point_msg.z = float(area)

            pub.publish(point_msg)
            rospy.loginfo(f"Published: x={point_msg.x}, y={point_msg.y}, id={point_msg.z}")
            
        cv2.imshow("Detected Red Block", frame_drawn)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        rate.sleep()

    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    try:
        red_block_publisher()
    except rospy.ROSInterruptException:
        pass